package com.zhouyu.service;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.builder.ExcelWriterBuilder;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.zhouyu.domain.User;
import com.zhouyu.entity.DemoData;
import com.zhouyu.mapper.UserMapper;
import com.zhouyu.utils.TestFileUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * @author lipengze
 * @version 1.0
 * @description: TODO
 * @date 2024/12/7 16:04
 */
@Service
@Slf4j
public class PoolZipTwoService {

    @Autowired
    private UserMapper userMapper;
    //设置线程池getSize() * 2，因为CPU密集型任务，线程数越多越好，可以提高处理效率
    private static final int THREAD_COUNT = Runtime.getRuntime().availableProcessors() * 2; // 根据CPU核心数设置线程池大小
    private static final int PAGE_SIZE = 500000; // 每个Excel的数据量
    private static final int TOTAL_EXCELS = 10; // 总共要导出的Excel数量
    private static final int BATCH_SIZE = 10000; // 每次查询的数据量

    /**
     * 导出多个Excel文件并压缩为ZIP包
     * @return 压缩后的ZIP文件的输入流
     * @throws IOException 如果文件操作过程中出现异常
     */
    public FileInputStream exportToZip() throws IOException {
        long startTime = System.currentTimeMillis(); // 记录开始时间

        List<String> excelFileNames = new ArrayList<>(); // 存储所有生成的Excel文件路径
        ExecutorService executorService = Executors.newWorkStealingPool(THREAD_COUNT); // 使用WorkStealingPool
        CountDownLatch countDownLatch = new CountDownLatch(TOTAL_EXCELS); // 创建CountDownLatch对象，用于等待所有线程完成

        for (int i = 1; i <= TOTAL_EXCELS; i++) { // 遍历所有Excel
            int finalI = i; // 因为lambda表达式需要的是有效的局部变量，所以这里创建一个final变量
            executorService.submit(() -> { // 提交任务到线程池
                try {
                    long sheetStartTime = System.currentTimeMillis(); // 记录单个Excel文件的开始时间
                    log.info("正在处理第{}个Excel", finalI); // 日志：正在处理的Excel编号

                    // 生成Excel文件的路径
                    String excelFileName = TestFileUtil.getPath() + "excel_" + finalI + "_" + System.currentTimeMillis() + ".xlsx";
                    ExcelWriterBuilder excelWriterBuilder = EasyExcel.write(excelFileName, DemoData.class);
                    WriteSheet writeSheet = EasyExcel.writerSheet("Sheet1").build(); // 创建WriteSheet对象

                    // 异步获取数据并写入Excel
                    fetchDataAndWriteToExcel(finalI, excelWriterBuilder, writeSheet);

                    long sheetEndTime = System.currentTimeMillis(); // 记录单个Excel文件的结束时间
                    log.info("第{}个Excel写入完成，耗时: {} 毫秒", finalI, sheetEndTime - sheetStartTime); // 输出单个Excel的写入耗时

                    synchronized (excelFileNames) {
                        excelFileNames.add(excelFileName); // 将Excel文件路径添加到列表中
                    }
                } catch (Exception e) {
                    log.error("第{}个Excel导出失败", finalI, e); // 如果发生异常，记录错误日志
                } finally {
                    countDownLatch.countDown(); // 无论成功还是失败，都减少CountDownLatch计数
                }
            });
        }

        try {
            countDownLatch.await(); // 等待所有线程完成
        } catch (InterruptedException e) {
            log.warn("线程被中断", e); // 记录中断日志
            Thread.currentThread().interrupt(); // 恢复中断状态
        }

        executorService.shutdown(); // 关闭线程池

        long endTime = System.currentTimeMillis(); // 记录结束时间
        log.info("所有Excel写入完成，总耗时: {} 秒", (endTime - startTime) / 1000.0); // 输出所有Excel的总写入耗时

        // 压缩Excel文件到ZIP
        long zipStartTime = System.currentTimeMillis(); // 记录压缩开始时间
        File zipFile = createZipFile(excelFileNames, "exportedData" + System.currentTimeMillis() + ".zip");
        long zipEndTime = System.currentTimeMillis(); // 记录压缩结束时间
        log.info("压缩完成，压缩耗时: {} 秒", (zipEndTime - zipStartTime) / 1000.0); // 输出压缩耗时

        // 返回压缩文件的输入流
        return new FileInputStream(zipFile);
    }

    /**
     * 异步获取数据并写入Excel
     * @param excelIndex  Excel索引（从1开始）
     * @param excelWriterBuilder  ExcelWriterBuilder对象
     * @param writeSheet  WriteSheet对象
     */
    private void fetchDataAndWriteToExcel(int excelIndex, ExcelWriterBuilder excelWriterBuilder, WriteSheet writeSheet) {
        try (ExcelWriter excelWriter = excelWriterBuilder.build()) {
            long lastId = 0; // 记录上一次查询的最大ID
            int totalWritten = 0; // 记录已写入的数据量

            while (totalWritten < PAGE_SIZE) {
                // 查询从 lastId 开始的下一页数据
                IPage<User> pageResult = userMapper.selectPage(new Page<>(1, BATCH_SIZE),
                        Wrappers.<User>lambdaQuery().gt(User::getId, lastId));

                if (pageResult.getRecords().isEmpty()) {
                    break; // 如果没有更多数据，退出循环
                }

                List<DemoData> demoDataList = pageResult.getRecords().stream()
                        .map(user -> {
                            DemoData demoData = new DemoData();
                            BeanUtils.copyProperties(user, demoData); // 使用BeanUtils进行对象转换
                            return demoData;
                        })
                        .collect(Collectors.toList());

                excelWriter.write(demoDataList, writeSheet); // 将数据写入Excel

                totalWritten += demoDataList.size();
                lastId = pageResult.getRecords().get(pageResult.getRecords().size() - 1).getId(); // 更新 lastId 为当前记录的 ID

                log.info("第{}个Excel已写入 {} 条数据", excelIndex, totalWritten);
            }
        } catch (Exception e) {
            log.error("第{}个Excel写入失败", excelIndex, e); // 如果发生异常，记录错误日志
        }
    }

    /**
     * 创建ZIP文件
     * @param excelFileNames Excel文件的路径列表
     * @param zipFileName ZIP文件的名称
     * @return 压缩后的文件对象
     * @throws IOException 如果文件操作过程中出现异常
     */
    private File createZipFile(List<String> excelFileNames, String zipFileName) throws IOException {
        File file = new File(TestFileUtil.getPath() + zipFileName);
        try (FileOutputStream fos = new FileOutputStream(file);
             ZipOutputStream zos = new ZipOutputStream(fos)) {
            for (String excelFileName : excelFileNames) {
                ZipEntry entry = new ZipEntry(new File(excelFileName).getName()); // 创建ZIP条目
                zos.putNextEntry(entry); // 添加条目到ZIP文件
                try (FileInputStream fis = new FileInputStream(excelFileName)) {
                    byte[] buffer = new byte[1024]; // 缓冲区
                    int length;
                    while ((length = fis.read(buffer)) > 0) {
                        zos.write(buffer, 0, length); // 将Excel文件的内容写入ZIP条目
                    }
                }
                zos.closeEntry(); // 关闭当前条目
                Files.delete(Paths.get(excelFileName)); // 删除原始Excel文件
            }
        }
        return file;
    }
}